# -*- coding: UTF-8 -*-
'''
@Project ：push_rk
@File ：model.py
@IDE ：PyCharm
@Author ：苦瓜
@Date ：2025/10/13 8:36
@Note: Something beautiful is about to happen !
'''
from dataclasses import dataclass, field
from typing import Tuple, Dict, List, Any

import matplotlib.pyplot as plt
import numpy as np
from keras import utils, Model, Sequential, optimizers, losses, metrics
from keras.src.layers import Conv2D, Flatten, Dense, Dropout, MaxPooling2D
from keras.src.datasets import cifar100


# 1.	模型约束描述
# 1)	数据加载与预处理：
# 从Keras的CIFAR-100数据集中加载训练和测试数据。
# 从CIFAR-100数据集中精选的4个类别如下：
# 1.	类别编号0 - apple（苹果）
# 2.	类别编号3 - baby（婴儿）
# 3.	类别编号53 - snail（蜗牛）
# 4.	类别编号79 - telephone（电话）
# 选择逻辑：
# 1.	跨领域多样性：包含自然物体（水果）、生物（人类/动物）和人造物品
# 2.	视觉复杂度梯度：从简单的几何形态（电话）到复杂有机形态（婴儿面部）
# 3.	语义关联性：建立潜在联想关系（婴儿→蜗牛→苹果→电话）
# 4.	分类挑战性：保留大小/形态相近的困难对（如圆形物体苹果vs蜗牛壳）
# 这些类别既保持了CIFAR-100原有的细粒度分类特性（20个超类中的4个不同超类），又能有效测试模型在跨域视觉概念理解上的能力。
#
# 加载后的数据被重新整形为适合CNN输入的格式（32x32x3的RGB图像），并将像素值归一化到0-1范围内。

@dataclass
class DataHandle:

    module : Any = None  # 必填
    filter_list : List[int] = None

    _x_train : np.ndarray = field(init=False, default=None)
    _y_train : np.ndarray = field(init=False, default=None)
    _x_test : np.ndarray = field(init=False, default=None)
    _y_test : np.ndarray = field(init=False, default=None)
    _y_cls_num : int = field(init=False, default=None)

    def __post_init__(self):
        if not self.module:
            raise ValueError("module 不能为空！")
        if not self.filter_list:
            # 设置默认类别
            self.filter_list = [0, 3, 53, 79]
        self._x_train, self._x_test, self._y_train, self._y_test = self.data_above_handle()
        ...

    def data_above_handle(self):
        """ 数据预处理 """
        # 加载数据
        (x_train, y_train),(x_test, y_test) = self.module.load_data()
        # 设置筛选条件
        train_filter = np.isin(y_train.flatten(), self.filter_list)
        test_filter = np.isin(y_test.flatten(), self.filter_list)
        # 筛选数据
        x_train = x_train[train_filter]
        y_train = y_train[train_filter]
        x_test = x_test[test_filter]
        y_test = y_test[test_filter]
        # 调整数据的维度，并归一化
        x_train, x_test = x_train.reshape(-1, 32, 32, 3) / 255. , x_test.reshape(-1, 32, 32, 3) / 255.
        # 映射类别
        y_train = self.map_class_dict(y_train)
        y_test = self.map_class_dict(y_test)
        print(f"""
        ---- 验证映射后类别数组 ----：
        {y_train}
        ---
        {y_test}
        ---- end ----
        """)
        # 类别数量
        self._y_cls_num = len(np.unique(y_train))
        # 对标签进行独热
        y_train = self.one_hot_code(y_train, self._y_cls_num)
        y_test = self.one_hot_code(y_test, self._y_cls_num)
        return x_train, x_test, y_train, y_test


    def one_hot_code(self, y, cls_num):
        return utils.to_categorical(y, cls_num)


    def map_class_dict(self, labels:np.ndarray) -> np.ndarray:
        """ 映射类别 """
        map_cls_dict = {
            old_cls : new_cls
            for new_cls, old_cls in enumerate(self.filter_list)
        }
        return np.vectorize(map_cls_dict.get)(labels)

    @property
    def return_train_test(self):
        return self._x_train, self._x_test, self._y_train, self._y_test


# 2)	模型定义：

class MyLeNet5(Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 卷积层
        self.conv = Sequential([
            Conv2D(filters=6, kernel_size=(5,5), activation='relu'),
            MaxPooling2D(),
            Conv2D(filters=16, kernel_size=(5, 5), activation='relu'),
            MaxPooling2D()
        ])
        # 展平层
        self.flatten = Sequential([
            Flatten()
        ])
        # 全连接层
        self.dense = Sequential([
            Dense(units=120, activation='relu'),
            Dropout(0.2),
            Dense(units=84, activation='relu'),
            Dropout(0.2),
            Dense(units=4, activation='softmax')
        ])


    def call(self, inputs, *args, **kwargs):
        x = self.conv(inputs)
        x = self.flatten(x)
        x = self.dense(x)
        return x


@dataclass
class ModelTrainer:

    data_handle : DataHandle = None  # 数据
    model : MyLeNet5 = None

    _history : Any = field(init=False, default=None)

    def __post_init__(self):
        if not self.model:
            raise ValueError("Trainer model 不能为空!")

        # 编译模型、训练模型
        self.model_compile_fit()

    def model_compile_fit(self) -> None:
        # 加载数据
        x_train, x_test, y_train, y_test = self.data_handle.return_train_test
        # 实例化模型
        self.model.build(input_shape=x_train.shape)
        self.model.compile(
            optimizer=optimizers.Adam(learning_rate=0.001),
            loss=losses.categorical_crossentropy,
            metrics=["acc", metrics.Recall, metrics.Precision, metrics.F1Score]
        )
        # 训练模型
        self._history = self.model.fit(
            x_train, y_train, batch_size=64,
            epochs=10, validation_data=(x_test, y_test)
        )

    @property
    def return_model_history(self):
        return self.model, self._history



# 3)	模型编译：
# 使用Adam优化器、稀疏分类交叉熵损失函数和准确率作为评估指标来编译模型。
# 4)	模型训练：
# 模型使用训练数据进行训练。
# 在训练过程中，使用验证数据来监控模型的性能。
# 5)	模型优化：
# 选择一种优化方法对模型进行进一步优化（如调整学习率、使用正则化、剪枝等），并分析优化后的结果。
# 6)	结果可视化：
# 训练完成后，提取训练和验证过程中的准确率数据。
# 使用Matplotlib库绘制训练和验证准确率的曲线图，以便可视化模型的训练过程。
# 7)	模型保存：
# 训练完成后，将模型的权重保存到文件中，以便将来可以重新加载和使用。
# 模型调优，考题中要求使用体现不同优化方案的要求，比如超参数方面的，比如说分类精度低的数据集的增强；比如说调整网络结构，比如说调整优化器学习率等；不能只是增加epoch大小，或者调整batch_size就可以，必须采用两种或者两种以上的调优方法进行调优，最终结论要给出综合各种方法调优得出的综合结论

def main():
    # 实例化数据
    data_handle = DataHandle(cifar100)
    # 实例化模型
    model_base = MyLeNet5()
    # 训练模型
    trainer = ModelTrainer(data_handle, model_base)
    # 获取到模型的信息与训练历史
    model, history = trainer.return_model_history
    # 打印模型形状
    model.summary()

    # 可视化模型准确率
    train_acc = history.history['acc']
    val_acc = history.history['val_acc']
    print(f"""
    模型的准确率
    acc: {train_acc}
    val_acc: {val_acc}
    """)
    plt.plot(train_acc, label="train acc", c="r")
    plt.plot(val_acc, label="val acc", c="g")
    plt.title("acc plot")
    plt.legend()
    plt.show()

    # 保存模型权重
    model.save_weights("model.weights.h5")
    # 加载模型权重进行评估
    model.load_weights("model.weights.h5")
    x_train, x_test, y_train, y_test = data_handle.return_train_test
    print(model.evaluate(x_test, y_test, return_dict=True))

"""
模型调优：
    当优化器为Adam：   - acc: 0.8379 - f1_score: 0.8361 - loss: 0.4087 - precision: 0.8551 - recall: 0.8198 - val_acc: 0.7950 - val_f1_score: 0.7923 - val_loss: 0.5136 - val_precision: 0.8057 - val_recall: 0.7775
    当优化器为Adagrad： acc: 0.4423 - f1_score: 0.3936 - loss: 1.3517 - precision: 0.0000e+00 - recall: 0.0000e+00 - val_acc: 0.4300 - val_f1_score: 0.3924 - val_loss: 1.3465 - val_precision: 0.0000e+00 - val_recall: 0.0000e+00
    当优化器为RMSprop： acc: 0.7814 - f1_score: 0.7801 - loss: 0.5084 - precision: 0.8085 - recall: 0.7536 - val_acc: 0.7225 - val_f1_score: 0.7191 - val_loss: 0.6712 - val_precision: 0.7598 - val_recall: 0.6800
    
    当优化器为Adam：
        学习率为 0.001 时： - acc: 0.8384 - f1_score: 0.8372 - loss: 0.4155 - precision: 0.8582 - recall: 0.8258 - val_acc: 0.8500 - val_f1_score: 0.8496 - val_loss: 0.4077 - val_precision: 0.8700 - val_recall: 0.8200
        学习率为 0.005 时：  - acc: 0.8082 - f1_score: 0.8069 - loss: 0.4582 - precision: 0.8197 - recall: 0.7932 - val_acc: 0.7575 - val_f1_score: 0.7577 - val_loss: 0.5318 - val_precision: 0.7778 - val_recall: 0.7350
        学习率为 0.01 时：  acc: 0.7867 - f1_score: 0.7833 - loss: 0.4919 - precision: 0.8111 - recall: 0.7606 - val_acc: 0.7925 - val_f1_score: 0.7915 - val_loss: 0.4873 - val_precision: 0.7959 - val_recall: 0.7700
        
        当学习率为0.001时，添加失活层：
            Dropout(0.2) : - acc: 0.8405 - f1_score: 0.8402 - loss: 0.3811 - precision: 0.8545 - recall: 0.8206 - val_acc: 0.8175 - val_f1_score: 0.8173 - val_loss: 0.4769 - val_precision: 0.8404 - val_recall: 0.7900
在优化器为Adam，学习率为0.001时，模型为当前最优。
"""


if __name__ == '__main__':
    main()